#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import sys
# sys.path.append('/opt/work/caffe/python')
sys.path.insert(0, '.')

import os
import json
import datetime
# from pycocotools.coco import COCO
import jpeg4py as jpeg
import shutil

# base_image_path = r'/rootfs/media/yery/Kaso/data/BDD100K/Images/100k/val'
# label_file_path = r'/rootfs/media/yery/Kaso/data/BDD100K/Labels/bdd100k_labels_images_val.json'
# out_image_path = r'/rootfs/media/yery/Kaso/data/BDD100K/out/Images/100k/val'
# out_coco_path = r'/rootfs/media/yery/Kaso/data/BDD100K/out/bdd100k_coco_labels_images_val.json'

base_image_path = r'/rootfs/media/yery/Kaso/data/BDD100K/Images/100k/train'
label_file_path = r'/rootfs/media/yery/Kaso/data/BDD100K/Labels/bdd100k_labels_images_train.json'
out_image_path = r'/rootfs/media/yery/Kaso/data/BDD100K/out/Images/100k/train'
out_coco_path = r'/rootfs/media/yery/Kaso/data/BDD100K/out/bdd100k_coco_labels_images_train.json'

# base_image_path = r'/rootfs/media/kasim/Data/data/BDD100K/Images/100k/val'
# label_file_path = r'/rootfs/media/kasim/Data/data/BDD100K/Labels/bdd100k_labels_images_val.json'
# out_image_path = r'/rootfs/media/kasim/Data/data/BDD100K/out/Images/100k/val'
# out_coco_path = r'/rootfs/media/kasim/Data/data/BDD100K/out/bdd100k_coco_labels_images_val.json'

# base_image_path = r'/rootfs/media/kasim/Data/data/BDD100K/Images/100k/train'
# label_file_path = r'/rootfs/media/kasim/Data/data/BDD100K/Labels/bdd100k_labels_images_train.json'
# out_image_path = r'/rootfs/media/kasim/Data/data/BDD100K/out/Images/100k/train'
# out_coco_path = r'/rootfs/media/kasim/Data/data/BDD100K/out/bdd100k_coco_labels_images_train.json'

# Train:
# person 95866
# cat 0
# dog 0
# baby_carriage 0
# face 0
# cell phone 0
# bicycle 7210
# motorcycle 3002

# Val
# person 13911
# cat 0
# dog 0
# baby_carriage 0
# face 0
# cell phone 0
# bicycle 1007
# motorcycle 452

category_set = {'motor', 'bike'}
coco_to_category = {
    'person': ['person', 'rider'],
    'motorcycle': 'motor',
    'bicycle': 'bike',
}

IMAGE_ID = 1000
BOX_OFFSET = 1000

CLASSES = ['person', 'cat', 'dog', 'baby_carriage', 'face', 'cell phone', 'bicycle', 'motorcycle']  # COCO

INFO = {
    "description": "Erised Dataset",
    "url": "http://www.erised.com",
    "version": "0.1.0",
    "year": 2020,
    "contributor": "Erised",
    "date_created": datetime.datetime.utcnow().isoformat(' ')
}

LICENSES = [
    {
        "id": 1,
        "name": "Attribution-NonCommercial-ShareAlike License",
        "url": "http://creativecommons.org/licenses/by-nc-sa/2.0/"
    }
]

CATEGORIES = [
    {
        'id': 1,
        'name': 'person',
        'supercategory': 'person',  # person
    },
    {
        'id': 17,
        'name': 'cat',
        'supercategory': 'cat',  # animal
    },
    {
        'id': 18,
        'name': 'dog',
        'supercategory': 'dog',  # animal
    },
    {
        'id': 102,
        'name': 'baby_carriage',
        'supercategory': 'baby_carriage',
    },
    {
        'id': 103,
        'name': 'face',
        'supercategory': 'face',  # person
    },
    {
        'id': 77,
        'name': 'cell phone',
        'supercategory': 'cell phone',  # electronic
    },
    {
        'id': 2,
        'name': 'bicycle',
        'supercategory': 'bicycle',  # vehicle
    },
    {
        'id': 4,
        'name': 'motorcycle',
        'supercategory': 'motorcycle',  # vehicle
    },
]

coco_label_ids = {}
coco_ids_label = {}
for category in CATEGORIES:
    category_name = category['name']
    category_id = category['id']
    label = coco_to_category.get(category_name, None)
    if label is None:
        continue
    if isinstance(label, (list, tuple)):
        for _label in label:
            coco_label_ids[_label] = category_id
    else:
        coco_label_ids[label] = category_id
    coco_ids_label[category_id] = category_name


def create_image_info(image_id, file_name, image_size,
                      date_captured=datetime.datetime.utcnow().isoformat(' '),
                      license_id=1, coco_url="", flickr_url=""):

    image_info = {
            "id": image_id,
            "file_name": file_name,
            "width": image_size[0],
            "height": image_size[1],
            "date_captured": date_captured,
            "license": license_id,
            "coco_url": coco_url,
            "flickr_url": flickr_url
    }

    return image_info


# annotation_id 也要唯一
def create_annotation_info(annotation_id, image_id, category_id, bounding_box=None, is_crowd=0):
    if bounding_box is None:
        return None

    # x1, y1, w, h = ann['bbox']
    area = bounding_box[2]*bounding_box[3]
    # if (area < MIN_BBOX_AREA) or (bounding_box[2] < MIN_BBOX_WIDTH) or (bounding_box[3] < MIN_BBOX_HEIGHT):
    #     return None

    annotation_info = {
        "id": annotation_id,
        "image_id": image_id,
        "category_id": category_id,
        "iscrowd": is_crowd,
        "area": area,
        "bbox": bounding_box,
        "segmentation": [],
    }

    return annotation_info


def main():
    if not os.path.exists(out_image_path):
        os.makedirs(out_image_path)
        os.system('chmod a+wr {}'.format(out_image_path))
    out_dir = os.path.dirname(out_coco_path)
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
        os.system('chmod a+wr {}'.format(out_dir))

    with open(label_file_path, 'r') as file:
        json_infos = json.load(file)
        # print(json_infos)

    coco_output = {
        "info": INFO,
        "licenses": LICENSES,
        "categories": CATEGORIES,
        "images": [],
        "annotations": []
    }

    obj_counts = {name: 0 for name in CLASSES}

    image_id = IMAGE_ID
    images = coco_output['images']
    annotations = coco_output['annotations']
    file_count = 0
    for json_info in json_infos:
        image_name = json_info['name']
        label_infos = json_info['labels']
        image_path = os.path.join(base_image_path, image_name)
        if not os.path.exists(image_path):
            continue
        box_id = 0
        has_bbox = False
        _annotations = []
        for label_info in label_infos:
            # {'box2d': {'x2': 129.76642, 'y1': 344.380116, 'y2': 400.529049, 'x1': 47.414654}, 'attributes': {'trafficLightColor': 'none', 'truncated': False, 'occluded': True}, 'manualShape': True, 'manualAttributes': True, 'id': 10, 'category': 'car'}
            category = label_info.get('category', None)
            if category is None:
                continue
            category_id = coco_label_ids.get(category, None)
            if category_id is None:
                continue
            category_id = int(category_id)
            bbox = label_info['box2d']
            bbox = [bbox['x1'], bbox['y1'], bbox['x2'], bbox['y2']]
            bounding_box = [bbox[0], bbox[1], (bbox[2] - bbox[0]), (bbox[3] - bbox[1])]

            obj_counts[coco_ids_label[category_id]] += 1
            annotation_id = box_id + image_id * BOX_OFFSET
            annotation_info = create_annotation_info(annotation_id, image_id, category_id, bounding_box)
            if annotation_info is None:
                continue
            _annotations.append(annotation_info)
            if category in category_set:
                has_bbox = True
            box_id += 1
        if has_bbox:
            annotations.extend(_annotations)
            jpg_file = jpeg.JPEG(image_path)
            jpg_file.parse_header()
            height = jpg_file.height
            width = jpg_file.width
            # width, height = 100, 100
            image_info = create_image_info(image_id, image_name, [width, height])
            images.append(image_info)
            image_id += 1
            dst_image_path = os.path.join(out_image_path, image_name)
            dst_image_dir = os.path.dirname(dst_image_path)
            if not os.path.exists(dst_image_dir):
                os.makedirs(dst_image_dir)
            shutil.copy(image_path, dst_image_path)
            file_count += 1
            if file_count % 100 == 0:
                print('Proccess File Count:', file_count)

    print('Proccess File Count:', file_count)
    with open(out_coco_path, 'w') as output_json_file:
        json.dump(coco_output, output_json_file)
    print('Save ', out_coco_path, 'Finish, File Count', file_count)
    os.system('chmod a+wr {}'.format(out_coco_path))
    os.system('chmod a+wr {} -R'.format(out_dir))

    print('=' * 80)
    for name in CLASSES:
        count = obj_counts[name]
        print(name, count)
    print('=' * 80)
    print('finish!')


if __name__ == '__main__':
    main()



